#!/bin/sh
#
# Patch or restore Ubuntu livecd iso file.
# Change content of file "grub.cfg". Set file size.

printHelp() {
	echo "Usage: sh %0 [-r] [-k] [-d] [-h|--help] [-v] [-p persistent-path] [-s resume-parameter] <iso-file>"
	echo "Options:"
	echo "  -r    Restore to original."
	echo "  -k    Keep grub.cfg file size."
	echo "  -d    Add debug parameters"
	echo "  -n    Add nomodeset parameter"
	echo "  -v    Verbose"
	echo "  -p persistent-path"
	echo "  -s resume-parameter"
	echo "  -h,--help    Get this help information."
	echo "  iso-file     ISO file to be patched."
	echo ""
	echo "  CAUTION:"
	echo "           Please backup your iso file before patch it."
	echo ""
	echo "  Example:"
	echo "    sh $0 -r lubuntu-22.04.5-desktop-amd64.iso"
	echo "    sh $0 -p lubuntu22045 -s \"resume=UUID=cb198d63-b5ca-4d95-a274-41fe82e417a8 resume_offset=659456\" lubuntu-22.04.5-desktop-amd64.iso"
}

# control vars
doKeepSize=false
doRestore=false
doDebug=false
doNomodeset=false
doVerbose=false
rs=""
pp=""

# parse arguments
while [ $# -gt 0 ]; do
case ${1#-} in
	r) doRestore=true; shift;;
	k) doKeepSize=true; shift;;
	d) doDebug=true; shift;;
	n) doNomodeset=true; shift;;
	v) doVerbose=true; shift;;
	p) shift; pp="$1"; shift;;
	s) shift; rs="$1"; shift;;
	h) printHelp; exit;;
	-help) printHelp; exit;;
	*)
		if [ x"$1" != x"${1#-}" ]; then
			echo "Unknown option: $1"
			printHelp
			exit 1
		else
			break
		fi
		shift;;
esac
done

# vars
n="${1%.iso}"
f="$n".iso; h="$n".iso.old.head
[ -n "$n" ] || { printHelp; exit 21; }
[ -e "$f" ] || { echo "Not found file: $f"; printHelp; exit 22; };

$doRestore && {
	# restore head part
	[ -s "$h" ] || { echo "Not found file: $h"; exit 31; }
	# shellcheck disable=SC2046
	{ dd if="$h" of="$f" bs=1M count=5 conv=notrunc $($doVerbose || echo status=none) &&
	cmp -n 4M "$h" "$f"; } || { echo "Restore failed."; exit 33; }
	exit
}

# backup old head part of livecd iso file
[ -e "$h" ] || {
	# shellcheck disable=SC2046
	dd if="$f" of="$h" bs=1M count=5 $($doVerbose || echo status=none) || { echo "Failed to backup to file $h"; exit 41; }
	$doVerbose && echo "Head part of iso file backup to file: $h"
}

# patch Ubuntu livecd iso file
# allow writable-backing and casper-sn on fs more than vfat, patch to delete "/root/$DEBCONF_TMPDIR", allow resume from hibernate
#
dl=/tmp/dd.$$.log
#
{ o=$(grep -m1 -h -b -U -z -P -e '^set timeout=30\n\nloadfont unicode\n\n([\s\S]*?)?\n\s+linux\s/casper/vmlinuz\s' "$h" | tr '\0' '\n') &&
	b=${o%%:*} && [ -n "$b" ] && b=$((b + 0)) && fs=$((${#o} - ${#b})); } || { echo "Failed to get grub.cfg file size."; exit 51; }
[ "$b" -ge $(( 5 * 1024 * 1024 * 1024 - 2048 )) ] && { echo "Content of grub.cfg locates out of $h"; exit 3; }
$doVerbose && echo "grub.cfg file size: $fs"
$doKeepSize || fs=2048
#
knh='linux /casper/vmlinuz'\
"$($doNomodeset && echo " nomodeset")"\
' noprompt noeject persistent'\
"$([ -n "$pp" ] && echo " persistent-path=$pp")"\
"$([ -n "$rs" ] && echo " \$rs")"
#
$doDebug && knh="$knh console=tty1 console=ttyS0,115200 ignore_loglevel keep break=mount"
#
# shellcheck disable=SC2016
knp='root=ram rdinit=/bin/sh -- -'"$($doDebug && echo xv)"'c "echo Run rdinit: \$0 \$@; p=\$1;r=\$2;a=\$3;b=\$4;t=\$5;c=/scripts/casper;f=\$c-helpers;v=false;dt=DEBCONF_TMPDIR; [ z\$- = z\${-#*x} ] || { v=true; env; doGrep(){ grep -n -H -A5 -B5 -e ^\\s\\+\$a\\\${p}\$b\\\$r\\?\$r -e ^\\s\\+\$t -e \\s\\+rm\s-rf\\s.*\$dt -e ^\\.\\s\\/scripts\\/local /init \$c \$f;}; doGrep;sleep 1;}; sed -i -e /^\\s\\+\$a\\\${p}\$b\\\$r\\?\$r\$/{s/\$a.*\$/\$t/} \$f; sed -i -e /^\\.\\s\\/scripts\\/local\$/{s/\$/\\nlocal_premount/} /init; sed -i -e /^\\s\\+rm\\s-rf\\s.*\$dt.\$/{s@rm@chroot\\t/root\\trm@} \$c; \$v && { sleep 1;doGrep;sleep 1;exec sh -xv /init;}; exec /init" rdinit \( \) vfat \|ext2 "vfat|ext2|ext4)"'
# shellcheck disable=SC2016
knps='root=ram rdinit=/bin/sh -- -'"$($doDebug && echo xv)"'c "echo \$0 \$@;p=\$1;r=\$2;a=\$3;b=\$4;t=\$5;c=/scripts/casper;f=\$c-helpers;dt=DEBCONF_TMPDIR; sed -i -e /^\\s\\+\$a\\\${p}\$b\\\$r\\?\$r\$/{s/\$a.*\$/\$t/} \$f; sed -i -e /^\\.\\s\\/scripts\\/local\$/{s/\$/\\nlocal_premount/} /init; sed -i -e /^\\s\\+rm\\s-rf\\s.*\$dt.\$/{s@rm@chroot\\t/root\\trm@} \$c; exec /init" rdinit \( \) vfat \|ext2 "vfat|ext2|ext4)"'
#
$doKeepSize && kn="$knh $knps" || kn="$knh $knp"
#
{ printf "%s\n" \
	'set timeout=5' \
	"$([ ${#kn} -lt $((fs - 100 - ${#rs})) ] && echo loadfont unicode)" \
	'menuentry "persist" {' \
	"$([ ${#kn} -lt $((fs - 90 - ${#rs})) ] && echo set gfxpayload=keep; \
	[ -n "$rs" ] && echo set rs=\""$rs"\")" \
	"$kn" \
	'initrd /casper/initrd' \
	'}' '' \
	"$([ ${#kn} -lt $((fs - 620)) ] && echo menuentry "${o#*menuentry }")" |
dd seek=$b oflag=seek_bytes conv=notrunc of="$f" 2>&1 | { if $doVerbose; then tee $dl; else cat > $dl; fi; } &&
grep -q "^[0-9]\+ bytes .*copied" $dl; } || { echo "Failed to write content of grub.cfg"; exit 55; }
#
sz=$(grep "^[0-9]\+ bytes .*copied" $dl | cut -d " " -f 1)
{ [ -n "$sz" ] && sz=$((sz + 0)); } || { echo "Failed to get file size."; exit 56; }
# fill remains
{ $doKeepSize && [ $sz -lt $fs ]; } && {
	$doVerbose && echo "Padding grub.cfg content remains."
	# shellcheck disable=SC2046
	dd seek=$((b + sz)) oflag=seek_bytes conv=notrunc of="$f" $($doVerbose || echo status=none) if=/dev/zero bs=$((fs - sz)) count=1 ||
		{ echo "Faild to padding remains."; exit 57; }
}
#
$doVerbose && echo "grub.cfg content changed to:"
# shellcheck disable=SC2046
grep $($doVerbose || echo "-q") -m1 -U -z -P -e '^set timeout=5\n' "$f" || { echo "Failed to change grub.cfg"; exit 58; }
#
#sz=$(grep "^[0-9]\+ bytes .*copied" $dl | cut -d " " -f 1)
#{ [ -n "$sz" ] && sz=$((sz + 0)); } || { echo "Failed to get file size."; exit 57; }
$doVerbose && echo "File content bytes written: $sz"
[ $fs -lt $sz ] && { echo "Written $sz bytes, exceeds file size $fs."; exit 59; }
#
$doKeepSize && { rm $dl; exit; }

# patch iso to set file size of "grub.cfg"
#[ -e "$f" ] && [ -e "$h" ] || exit $?
#
{ o=$(grep -m1 -bUz -P -e '^\x01\x0aGRUB\.CFG\;1$' "$h" | tr '\0' '\n') &&
a=${o%%:*} && true echo "len(o - a - 1)=$((${#o} - ${#a} -1))" &&
a=$((a - 21)) && [ $a -gt 0 ]; } || { echo "Failed to get file size offset."; exit 63; }
[ $a -ge $(( 5 * 1024 * 1024 * 1024 - 128 )) ] && { echo "grub.cfg file size locates out of $h"; exit 3; }
$doVerbose && {
	echo "File size offset: $a"
	od -Ad -j$a -N48  -t x1a "$h"
}
# shellcheck disable=SC2046
printf "%b" "$(printf "%s%03o" \\ $((sz % 256)) \\ $((sz >> 8)) \\ 0 \\ 0 \\ 0 \\ 0 \\ $((sz >> 8)) \\ $((sz % 256)) )" |
dd seek=$a oflag=seek_bytes conv=notrunc of="$f" $($doVerbose || echo status=none) || { echo "Faild to change file size."; exit 67; }
#
$doVerbose && {
	echo "Change file size to: $(printf "0x%04x" "$fs")"
	od -Ad -j$((a - 24)) -N 144 -t x1a "$f"
}
#
rm $dl

